package com.micro.ai.commons.client;

import com.micro.ai.commons.domain.ApiResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.reactive.function.client.WebClientResponseException;

import java.time.Duration;
import java.util.Map;

/**
 * Ollama客户端
 * 
 * @author micro-ai
 * @since 0.0.1
 */
@Slf4j
@Component
public class OllamaClient {

    private final WebClient webClient;
    private final String baseUrl;

    public OllamaClient(WebClient webClient) {
        this.webClient = webClient;
        this.baseUrl = "http://localhost:11434"; // 默认Ollama服务地址
    }

    /**
     * 生成文本
     */
    public ApiResponse<String> generateText(String model, String prompt) {
        return generateText(model, prompt, null, null);
    }

    /**
     * 生成文本（带参数）
     */
    public ApiResponse<String> generateText(String model, String prompt, String systemPrompt, Map<String, Object> options) {
        try {
            GenerateRequest request = GenerateRequest.builder()
                    .model(model)
                    .prompt(prompt)
                    .system(systemPrompt)
                    .options(options)
                    .stream(false)
                    .build();

            GenerateResponse response = webClient.post()
                    .uri(baseUrl + "/api/generate")
                    .bodyValue(request)
                    .retrieve()
                    .bodyToMono(GenerateResponse.class)
                    .timeout(Duration.ofSeconds(30))
                    .block();

            if (response != null && response.response() != null) {
                return ApiResponse.<String>success(response.response());
            } else {
                return ApiResponse.error("D0004", "AI推理失败");
            }
        } catch (WebClientResponseException e) {
            log.error("Ollama API调用失败: {}", e.getMessage(), e);
            return ApiResponse.error("D0009", "AI服务不可用: " + e.getMessage());
        } catch (Exception e) {
            log.error("Ollama客户端异常: {}", e.getMessage(), e);
            return ApiResponse.error("D0004", "AI推理失败: " + e.getMessage());
        }
    }

    /**
     * 批量生成文本
     */
    public ApiResponse<String> generateBatch(String model, String[] prompts) {
        StringBuilder result = new StringBuilder();
        
        for (String prompt : prompts) {
            ApiResponse<String> response = generateText(model, prompt);
            if (response.isSuccess()) {
                result.append(response.data()).append("\n");
            } else {
                return response;
            }
        }
        
        return ApiResponse.<String>success(result.toString());
    }

    /**
     * 健康检查
     */
    public ApiResponse<Boolean> healthCheck() {
        try {
            webClient.get()
                    .uri(baseUrl + "/api/tags")
                    .retrieve()
                    .bodyToMono(String.class)
                    .timeout(Duration.ofSeconds(5))
                    .block();
            
            return ApiResponse.success(true);
        } catch (Exception e) {
            log.warn("Ollama健康检查失败: {}", e.getMessage());
            return ApiResponse.success(false);
        }
    }

    /**
     * 生成请求模型
     */
    public record GenerateRequest(
        String model,
        String prompt,
        String system,
        Map<String, Object> options,
        Boolean stream
    ) {
        public static GenerateRequestBuilder builder() {
            return new GenerateRequestBuilder();
        }
        
        public static class GenerateRequestBuilder {
            private String model;
            private String prompt;
            private String system;
            private Map<String, Object> options;
            private Boolean stream;
            
            public GenerateRequestBuilder model(String model) {
                this.model = model;
                return this;
            }
            
            public GenerateRequestBuilder prompt(String prompt) {
                this.prompt = prompt;
                return this;
            }
            
            public GenerateRequestBuilder system(String system) {
                this.system = system;
                return this;
            }
            
            public GenerateRequestBuilder options(Map<String, Object> options) {
                this.options = options;
                return this;
            }
            
            public GenerateRequestBuilder stream(Boolean stream) {
                this.stream = stream;
                return this;
            }
            
            public GenerateRequest build() {
                return new GenerateRequest(model, prompt, system, options, stream);
            }
        }
    }

    /**
     * 生成响应模型
     */
    public record GenerateResponse(
        String response,
        Boolean done,
        String model,
        Long created_at,
        Long total_duration,
        Long load_duration,
        Long prompt_eval_duration,
        Long eval_duration,
        Integer prompt_eval_count,
        Integer eval_count
    ) {}
}
